# coding=utf-8
import time

from codeobject import CodeObject


class Parser(object):
    def __init__(self, file_path):
        self._pyc_file = open(file_path, 'r')
        self.magic_number = None
        self.modify_time = None
        self.code_object = None

        self._string_table = []

    def parse(self):
        self.magic_number = "0x{:08x}".format(self._read_int())
        # python2.7 的魔数是 0x0a0df303
        assert self.magic_number == "0x0a0df303"

        self.modify_time = time.ctime(self._read_int())

        # 根据规定, 后面的一个字符一定是 c, 表示代码段
        assert self._pyc_file.read(1) == 'c'

        return self._get_code_object()

    def _get_code_object(self):
        code_object = CodeObject()

        code_object.arg_count = self._read_int()
        code_object.nlocals = self._read_int()
        code_object.stack_size = self._read_int()
        code_object.flag = self._read_int()
        code_object.byte_codes = self._get_byte_code()
        code_object.consts = self._get_consts()
        code_object.names = self._get_names()
        code_object.var_names = self._get_var_names()
        code_object.free_vars = self._get_free_vars()
        code_object.cell_vars = self._get_var_names()
        code_object.file_name = self._get_file_name()
        code_object.co_name = self._get_name()
        code_object.line_no = self._read_int()
        code_object.lnotable = self._get_no_table()

        return code_object

    def _read_int(self):
        data = bytearray(self._read(4))
        m = data[0]
        m |= data[1] << 8
        m |= data[2] << 16
        m |= data[3] << 24
        return m

    def _get_byte_code(self):
        assert self._read(1) == 's'
        return self._get_string()

    def _get_consts(self):
        if self._read(1) == '(':
            return self._get_tuple()
        self._read_back(1)
        return []

    def _get_names(self):
        if self._read(1) == '(':
            return self._get_tuple()
        self._read_back(1)
        return []

    def _get_var_names(self):
        if self._read(1) == '(':
            return self._get_tuple()
        self._read_back(1)
        return []

    def _get_free_vars(self):
        if self._read(1) == '(':
            return self._get_tuple()
        self._read_back(1)
        return []

    def _get_cell_vars(self):
        if self._read(1) == '(':
            return self._get_tuple()
        self._read_back(1)
        return []

    def _get_file_name(self):
        return self._get_name()

    def _get_no_table(self):
        ch = self._read(1)
        if ch not in ['s', 't']:
            return None

        return self._get_string()

    def _get_tuple(self):
        length = self._read_int()
        l = []
        for _ in range(length):
            obj_type = self._read(1)

            if obj_type == 'c':
                l.append(self._get_code_object())
            elif obj_type == 'i':
                l.append(self._read_int())
            elif obj_type == 'N':
                l.append(None)
            elif obj_type == 't':
                s = self._get_string()
                l.append(s)
                self._string_table.append(s)
            elif obj_type == 's':
                l.append(self._get_string())
            elif obj_type == 'R':
                index = self._read_int()
                l.append(self._string_table[index])
            else:
                print("Parser Failed. Unknown Type :", obj_type)
        return l

    def _get_name(self):
        ch = self._read(1)

        if ch == 's':
            return self._get_string()
        elif ch == 't':
            s = self._get_string()
            self._string_table.append(s)
            return s
        elif ch == 'R':
            index = self._read_int()
            return self._string_table[index]
        else:
            return None

    def _get_string(self):
        length = self._read_int()
        data = self._read(length)
        return data

    def _read(self, n):
        return self._pyc_file.read(n)

    def _read_back(self, n):
        """
        将当前读取的游标回退n个字节
        :param n: 回退的个数
        :return:
        """
        # 第二个参数 1 代表从当前位置当做起始点
        self._pyc_file.seek(-n, 1)

    def close(self):
        self._pyc_file.close()


if __name__ == '__main__':
    pyc_file_path = 'test/simple_test.pyc'
    parser = Parser(pyc_file_path)
    c = parser.parse()
    parser.close()
    print(parser.magic_number, parser.modify_time)
